package main

import (
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"

	"golang.org/x/net/dns/dnsmessage"
)

func main() {
	var err error

	//加载ESNI密钥
	var esniKeys *tls.ESNIKeys
    var esniPrivateKey []byte
	esniPrivateKey, err = ioutil.ReadFile("esni")
	if err != nil {
		log.Fatalf("Failed to read ESNI private key: %s", err)
	}
	contents, err := ioutil.ReadFile("esni.pub")
	if err != nil {
		log.Fatalf("Failed to read ESNIKeys: %s", err)
	}
	esniKeysBytes, err := base64.StdEncoding.DecodeString(string(contents))
	if err != nil {
		log.Fatal("Bad -esni-keys: %s", err)
	}
	esniKeys, err = tls.ParseESNIKeys(esniKeysBytes)
	if esniKeys == nil {
		log.Fatalf("Cannot parse ESNIKeys: %s", err)
	}

	//配置TLS
	tlsConfig := &tls.Config{
		GetCertificate:getCertificateForSNI,
		MinVersion:       tls.VersionTLS13,
		MaxVersion:       tls.VersionTLS13,
		CurvePreferences: []tls.CurveID{},
		GetServerESNIKeys:func([]byte) (*tls.ESNIKeys, []byte, error) { return esniKeys, esniPrivateKey, nil },
		Accept0RTTData: true,
	}

	//创建一个http.server实例
	server := &http.Server{
		Addr: ":443",
		TLSConfig: tlsConfig,
	}

	// 注册 DoH 请求处理函数
	http.HandleFunc("/dns-query", handleESNIRequest)

	// 启动 HTTPS 服务器
	log.Println("Starting DoH server on https://localhost:443")
	if err := server.ListenAndServeTLS("", ""); err != nil {
		log.Fatal("ListenAndServeTLS: ", err)
	}


}

// 处理 ESNI 请求
func handleESNIRequest(w http.ResponseWriter, req *http.Request) {
	// 读取 ESNI 请求中的 DNS 消息
	dnsRequest := req.Body
	defer dnsRequest.Close()

	// 读取请求中的 DNS 消息内容
	var buf bytes.Buffer
	_, err := io.Copy(&buf, dnsRequest)
	if err != nil {
		log.Println("Failed to read DNS request: ", err)
		http.Error(w, "Failed to read DNS request", http.StatusInternalServerError)
		return
	}

	// 解包 DNS 消息
	var msg dnsmessage.Message
	err = msg.Unpack(buf.Bytes())
	if err != nil {
		log.Println("Failed to unpack DNS request: ", err)
		http.Error(w, "Failed to unpack DNS request", http.StatusInternalServerError)
		return
	}

	// 向真实的 DNS 服务器发送请求
	dnsResponse, err := sendDNSRequestToRealServer(msg)
	if err != nil {
		log.Println("Failed to send DNS request to real server: ", err)
		http.Error(w, "Failed to get DNS response", http.StatusInternalServerError)
		return
	}

	// 发送 DNS 响应作为 DoH 响应
	w.Header().Set("Content-Type", "application/dns-message")
	w.Write(dnsResponse)
}

// 向真实的 DNS 服务器发送 DNS 请求
func sendDNSRequestToRealServer(msg dnsmessage.Message) ([]byte, error) {
	// 连接到 DNS 服务器
	dnsServer := "114.114.114.114:53"
	conn, err := net.Dial("udp", dnsServer)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to DNS server: %w", err)
	}
	defer conn.Close()

	// 将 DNS 消息序列化
	dnsBytes, err := msg.Pack()
	if err != nil {
		return nil, fmt.Errorf("failed to pack DNS message: %w", err)
	}

	// 发送 DNS 请求
	_, err = conn.Write(dnsBytes)
	if err != nil {
		return nil, fmt.Errorf("failed to send DNS request: %w", err)
	}

	// 读取 DNS 响应
	responseBuf := make([]byte, 512)
	n, err := conn.Read(responseBuf)
	if err != nil {
		return nil, fmt.Errorf("failed to read DNS response: %w", err)
	}

	return responseBuf[:n], nil
}

func getCertificateForSNI(helloInfo *tls.ClientHelloInfo) (*tls.Certificate, error) {
	switch helloInfo.ServerName {
    case "server1.com":
        cert, err := tls.LoadX509KeyPair("server1.crt", "server1.key")
        return &cert, err  // 返回结构体的指针
    case "server2.com":
        cert, err := tls.LoadX509KeyPair("server2.crt", "server2.key")
        return &cert, err  // 返回结构体的指针
    default:
        return nil, fmt.Errorf("no certificate found for SNI: %s", helloInfo.ServerName)
    }
}
